from ctypes import *
import cv2, time, gc, string, queue, argparse, threading, random, os
import numpy as np 
class decoderParam(Structure):
  _fields_ = [    
    ("jump", c_int),
    ("lmtMB", c_int),
    ("rawDataType", c_int),   # //0:YUVI420  1:BGR_CHW  2:RGB_CHW   3:BGR_HWC  4:RGB_HWC
    ("retDecodeHeigth", c_int),
    ("retDecodeWidth", c_int),
    ("retFrameSize", c_int),
    ("retOutCnt", c_int),
    ("retFps", c_int),
    ("retMemFrameNum", c_int),
    ("cropRight", c_int),
    ("cropLeft", c_int),
    ("cropBottom", c_int),
    ("cropTop", c_int),
    ("resizeH", c_int),
    ("resizeW", c_int),  
    ("memWithTimestamp", c_int),    
    ("retMemTimestampOffset", c_int),    
    ("startTime", c_int64),
    ("retEndTime", c_int64),
        ]
  def show(self):
    print("***************decoder ret***************")
    print("\t* retDecodeWidth =",self.retDecodeWidth)
    print("\t* retDecodeHeigth =",self.retDecodeHeigth)
    print("\trawDataType =",self.rawDataType) 
    print("\tjump =",self.jump) 
    print("\tlmtMB =",self.lmtMB)  
    print("\tretFrameSize =",self.retFrameSize)
    print("\tretOutCnt =",self.retOutCnt)
    print("\tretFps =",self.retFps)
    print("\tretMemFrameNum =",self.retMemFrameNum)
    print("\tmemWithTimestamp =",self.memWithTimestamp)
    print("\tretMemTimestampOffset =",self.retMemTimestampOffset)
    print("\tstartTime =",self.startTime)
    print("\tretEndTime =",self.retEndTime)
    print("*****************************************")
def pyStr2C(pyStr):
    slen=len(pyStr)
    if slen>255:
        print("error! path length less than 255!")
        exit(0)
    return (c_char * 256)(*bytes(pyStr, 'utf-8'))
def YUVtoRGB(byteArray,width,height):
    e = width * int(height*1.5)
    Y = byteArray[0:e]
    Y = np.reshape(Y, (int(height*1.5), width))
    RGBMatrix = Y.astype(np.uint8)
    RGBMatrix = cv2.cvtColor(RGBMatrix, cv2.COLOR_YUV2RGB_NV21, 3)
    return RGBMatrix
dll = cdll.LoadLibrary('../build/unreDec.so');
#=================================================
DecoderMem = dll.DecoderMem2
DecoderMem.restype = c_int
DecoderMem.argtypes = [c_int,c_char_p,c_char_p,POINTER(decoderParam)]
#=================================================
DecoderFile = dll.DecoderFile
DecoderFile.restype = c_int
DecoderFile.argtypes = [c_int,c_char_p,c_char_p,POINTER(decoderParam)]
#=================================================
splitMp4 = dll.splitMp4
splitMp4.restype = c_int
splitMp4.argtypes = [c_char_p,c_char_p,POINTER(decoderParam)] 
#=================================================
DecoderMem3 = dll.DecoderMem3
DecoderMem3.restype = c_int
DecoderMem3.argtypes = [c_int,c_char_p,POINTER(decoderParam)]
#=================================================
DecoderOpen = dll.Open2
DecoderOpen.restype = c_int
DecoderOpen.argtypes = [c_int,c_char_p,POINTER(decoderParam)]
#=================================================
splitMp4C = dll.splitMp4C
splitMp4C.restype = c_int
splitMp4C.argtypes = [c_char_p,c_char_p,c_int] 
#=================================================

def splitList(listTemp, n):
    s=[]
    eachLength = (len(listTemp)+len(listTemp)%n)//n
    for i in range(0, len(listTemp), eachLength):
        b=listTemp[i:i+eachLength]
        s.append(b)
    if len(s)>n:s[n-1]=s[n-1]+s[n]
    s=s[0:n]
    return s


def unpack(mem, param):
    frameSize = param.retFrameSize
    timestampOffset = param.retMemTimestampOffset
    timestamps = np.frombuffer(mem[0:timestampOffset], np.int64)
    print("timestamps =",timestamps)
    # assert(len(timestamps) == param.retMemFrameNum),len(timestamps)
    img_np=None
    for i in range(param.retMemFrameNum):
        aFrame = mem[timestampOffset+frameSize*i:timestampOffset+frameSize*i+frameSize] 
        # print("thisTimestamp =",timestamps[i])
        nparr = np.frombuffer(aFrame, np.uint8)
        if param.rawDataType==0:
            img_np = YUVtoRGB(nparr,param.retDecodeWidth,param.retDecodeHeigth) 
        elif param.rawDataType==1 or param.rawDataType==2:
            img_np = np.reshape(nparr, (3,param.retDecodeHeigth, param.retDecodeWidth)) 
            #print(img_np.shape)
            img_np=img_np.transpose(1,2,0)
            #print(img_np.shape)
            #exit(0)
        elif param.rawDataType==3 or param.rawDataType==4:
            img_np = np.reshape(nparr, (param.retDecodeHeigth, param.retDecodeWidth,3)) 
        # cv2.imshow("123",img_np)
        # cv2.waitKey(0) 
        # cv2.imwrite(str(i)+'_'+str(timestamps[i])+'.png', img_np)
        # cv2.imwrite(str(timestamps[i])+'_'+str(i)+'.bmp', img_np)

def signThread(path):
    initRet = dll.DecoderInit2(0)
    print("************initRet =",initRet)
    param=decoderParam()
    param.startFrameIndex=0;
    #param.startFrameIndex=3452;
    param.jump=1
    param.lmtMB=600 #5GB
    param.rawDataType=3
    param.startTime=0
    param.memWithTimestamp=1
    fileIn = pyStr2C(path)
    mem = create_string_buffer(param.lmtMB*1024*1024)
    t0=time.time()
    decRet = DecoderMem(0,fileIn,mem,byref(param))
    assert(decRet==0),decRet
    #if decRet==0: param.show()    
    print('time='+str(time.time()-t0))    
    print("retEndTime =",param.retEndTime)
    destroyRet = dll.DecoderDestroy2()
    print("#########destroyRet =",destroyRet)
    unpack(mem, param)
    mem =None
    

def cpudec(path, n=100):
    cap=cv2.VideoCapture(path)
    frameindex = 0
    ret=True
    while ret and frameindex<n:
        ret,frame = cap.read()
        # cv2.imwrite(str(frameindex)+'_c.bmp', frame)
        frameindex += 1
    cap.release()
    print("frameindex=",frameindex)


def gpudec(path, idx, startIndex, endIndex):

    cap=cv2.VideoCapture(path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()

    print("fps=",fps)

    initRet = dll.DecoderInit2(idx)
    print("************initRet =",initRet)
    param=decoderParam()

    param.jump=1
    param.lmtMB=1024#5GB
    param.rawDataType=3
    param.memWithTimestamp=1
    
    fileIn = pyStr2C(path)



    param.retEndTime=0
    
    
    totalc=0
    mem = create_string_buffer(param.lmtMB*1024*1024)
    param.startTime=param.retEndTime
    decRet=DecoderOpen(idx, fileIn,byref(param))
    
    while param.retEndTime>=0 :        
        decRet = DecoderMem3(idx, mem,byref(param))
        unpack(mem, param)
        param.show()
        totalc+=param.retMemFrameNum
        
    destroyRet = dll.DecoderDestroy2()
    print("#########destroyRet =",destroyRet)
    print("totalc=",totalc, "idx=", idx)
    
    # totalframe = int(cv2.VideoCapture(path).get(cv2.CAP_PROP_FRAME_COUNT))
    # print("totalframe=",totalframe)



def split(path, lmtMB):
    t0=time.time()    
    initRet = dll.DecoderInit(0)
    print("************initRet =",initRet)
    param=decoderParam()
    param.lmtMB=lmtMB#5GB    
    ret=splitMp4(pyStr2C(path),pyStr2C(path),byref(param))
    print('ret=',ret)
    param.show()
    destroyRet = dll.DecoderDestroy()
    print("#########destroyRet =",destroyRet)
    print('time='+str(time.time()-t0))
    t0=time.time()    
    tt=0
    ret=[]
    for i in range(param.retOutCnt):
        cap=cv2.VideoCapture(path+str(i)+'.mp4')
        totalframe = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        # print('totalframe=',totalframe)
        cap.release()
        ret.append([path+str(i)+'.mp4', tt])
        tt+=totalframe
    print('tt=',tt)
    cap=cv2.VideoCapture(path)
    totalframe0 = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    print('totalframe0=',totalframe0)
    print('time='+str(time.time()-t0))
    return ret, tt

def split2(path, nSeg):
    t0=time.time()    
    name, ext=os.path.splitext(path)
    outtpl=name+'_%d'+ext
    retN=splitMp4C(pyStr2C(path),pyStr2C(outtpl),nSeg)
    print('retN=',retN)
    print('time='+str(time.time()-t0))
    t0=time.time()    
    tt=0
    ret=[]
    for i in range(retN):
        cpath=name+'_'+str(i)+ext
        cap=cv2.VideoCapture(cpath)
        totalframe = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        # print('totalframe=',totalframe)
        cap.release()
        ret.append([cpath, tt])
        tt+=totalframe
    print('tt=',tt)
    cap=cv2.VideoCapture(path)
    totalframe0 = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    print('totalframe0=',totalframe0)
    print('time='+str(time.time()-t0))
    return ret, tt


def gpudec1(queueOut, path, idx, startIndex, endIndex, param, mem):
    print('start gpudec1')
    fileIn = pyStr2C(path)
    param.startTime=0
    decRet=DecoderOpen(idx, fileIn, byref(param))
    # print('open gpudec1')
    ci=startIndex
    img_np=None
    while param.retEndTime>=0 :        
        decRet = DecoderMem3(idx, mem, byref(param))
        # print('DecoderMem3 gpudec1')
        frameSize = param.retFrameSize
        timestampOffset = param.retMemTimestampOffset
        timestamps = np.frombuffer(mem[0:timestampOffset], np.int64)
        # print("timestamps =",timestamps)
        # param.show()
        # assert(len(timestamps) == param.retMemFrameNum),len(timestamps)
        for i in range(param.retMemFrameNum):
            aFrame = mem[timestampOffset+frameSize*i:timestampOffset+frameSize*i+frameSize] 
            # print("thisTimestamp =",timestamps[i])
            nparr = np.frombuffer(aFrame, np.uint8)
            if param.rawDataType==0:
                img_np = YUVtoRGB(nparr,param.retDecodeWidth,param.retDecodeHeigth) 
            elif param.rawDataType==1 or param.rawDataType==2:
                img_np = np.reshape(nparr, (3,param.retDecodeHeigth, param.retDecodeWidth)).transpose(1,2,0)
            elif param.rawDataType==3 or param.rawDataType==4:
                img_np = np.reshape(nparr, (param.retDecodeHeigth, param.retDecodeWidth,3)) 
            if img_np is not None:
                queueOut.put([ [ci, '2020-20-20 20:20:20'], img_np.copy() ])
            else:
                queueOut.put([ [ci, '2020-20-20 20:20:20'], img_np ])
            ci+=1
    while ci<endIndex:
        if img_np is not None:
            queueOut.put([ [ci, '2020-20-20 20:20:20'], img_np.copy() ])
        else:
            queueOut.put([ [ci, '2020-20-20 20:20:20'], img_np ])
        ci+=1
    # queueOut.put(None)
    param.retEndTime=0
    print('end gpudec1')
    
def gpudeca(queueOut, tasklist, idx, endIndex):
    initRet = dll.DecoderInit2(idx)
    print("************initRet =",initRet)
    param=decoderParam()

    param.jump=1
    param.lmtMB=1024#5GB
    param.rawDataType=3
    param.memWithTimestamp=1
    mem = create_string_buffer(param.lmtMB*1024*1024)
    
    if len(tasklist)>1:
        for i in range(len(tasklist)-1):
            gpudec1(queueOut, tasklist[i][0], idx, tasklist[i][1], tasklist[i+1][1], param, mem)
            os.remove(tasklist[i][0])
    gpudec1(queueOut, tasklist[len(tasklist)-1][0], idx, tasklist[len(tasklist)-1][1], endIndex, param, mem)
    os.remove(tasklist[len(tasklist)-1][0])
    queueOut.put(None)
    # destroyRet = dll.DecoderDestroy2()
    # print("#########destroyRet =",destroyRet)

def Pass(queueIn, nStopIn=1):
    while nStopIn:
        info = queueIn.get()
        if info is None:
            nStopIn-=1


def jumpt(path, n=4):
    t0=time.time()    
    seglist, totalc=split(path, 1024)
    print(str(seglist))
    step=int(len(seglist)/n)+1
    qtmp = [ queue.Queue(2) for i in range(n) ]
    thl = [ threading.Thread(target=gpudeca, args=(qtmp[i], seglist[(i*step):((i+1)*step)], i, seglist[(i+1)*step][1])) for i in range(n-1) ]
    thl.extend( [ threading.Thread(target=Pass, args=(qtmp[i], 1)) for i in range(n) ] )
    for thi in thl:
        thi.start()
    gpudeca(qtmp[n-1], seglist[((n-1)*step):(len(seglist))], (n-1), totalc)
    for thi in thl:
        thi.join()
    destroyRet = dll.DecoderDestroy2()
    print("#########destroyRet =",destroyRet)
    et=time.time()-t0
    print('time='+str(et)+'\tfps='+str(totalc/et))





if __name__=='__main__':
    path="/home/u/20200702/HD_video1.mp4"
    # path="test.mp4"
    split2(path, 4)
    # jumpt(path, 2)
    
    # cap=cv2.VideoCapture(path)
    # fps = cap.get(cv2.CAP_PROP_FPS)
    # totalframe0 = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # print('totalframe0=',totalframe0)
    # cap.release()

    # signThread(path)
    # cpudec(path, 11010)
    # gpudec(path, 0, 0, 75)
    
    # m1=cv2.imread('./1/2068_2.bmp')
    # m2=cv2.imread('./2/60_c.bmp')
    
    # df=cv2.absdiff(m1,m2)
    # cv2.imshow('df',df)
    # cv2.waitKey()
    
    # t0=time.time()    
    # initRet = dll.DecoderInit(0)
    # print("************initRet =",initRet)
    # param=decoderParam()
    # param.lmtMB=1024#5GB    
    # ret=splitMp4(pyStr2C(path),pyStr2C(path),byref(param))
    # print('ret=',ret)
    # param.show()
    # destroyRet = dll.DecoderDestroy()
    # print("#########destroyRet =",destroyRet)
    # print('time='+str(time.time()-t0))
    # t0=time.time()    
    # tt=0
    # for i in range(param.retOutCnt):
        # cap=cv2.VideoCapture(path+str(i)+'.mp4')
        # totalframe = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        # # print('totalframe=',totalframe)
        # cap.release()
        # tt+=totalframe
    # print('tt=',tt)
    # print('totalframe0=',totalframe0)
    # print('time='+str(time.time()-t0))
    
    exit(0)
 
 

